last.fxn <- function(x) x[length(x)]

date.fxn <- function(start, end) {
  start.date <- as.Date(start, tz="UTC")
  end.date <- strptime(end, format="%Y-%m-%d", tz="UTC")
  exposure.duration <- as.numeric(difftime(as.POSIXct(end.date), as.POSIXct(start.date, tz="UTC"), units="days")) + 1 
  return(exposure.duration)
}

population.fxn <- function(world.population, country, year, starting.ages) {
  popl <- subset(world.population, Time==year & Location==country)
  popl$AgeCat <- cut(popl$AgeGrpStart, c(c(starting.ages-1), 120))
  popl.count <- as.numeric(by(1000*popl$PopTotal, popl$AgeCat, function(x) sum(x)))
  names(popl.count) <- starting.ages
  return(popl.count)
}

create.data.wide <- function(starting.ages, population, cases, deaths) {
  data.wide <- data.table(cbind(age=starting.ages, population=population, cases=cases, deaths=deaths))
  data.wide[, ("prevalence"):=cases/population]
  data.wide[, ("exposure"):=population]
  data.wide[, ("log.exposure"):=log(population)]
  data.wide[, ("case.fatality.rates"):=deaths/cases]
  data.wide[which(deaths==0),"case.fatality.rates"] <- 0
  return(data.wide)
}

create.data.long <- function(data.wide) {
  tmp <- data.wide
  tmp[,`:=`(w0= exposure)]	
  names(tmp)[which(names(tmp)=="deaths")]<-"w1"
  tmp[,id:=seq(1,dim(tmp)[1])]
  tmp2 <- data.table(reshape(tmp[,.(w1,w0,id,age,case.fatality.rates, prevalence)],idvar="id",varying=list(c(1:2)),direction="long",v.names="w") )
  tmp2[,deaths:=time]
  tmp2[case.fatality.rates>0,log.case.fatality.rates:=log(case.fatality.rates)]
  tmp2[case.fatality.rates==0,log.case.fatality.rates:=0]
  tmp2[,log.prevalence:=log(prevalence)]
  tmp2[time==2,deaths:=0]
  deaths.data.long <-tmp2; rm(tmp,tmp2)
  return(deaths.data.long)
}

covid19.mortality.rate.fxn <- function(starting.ages, population, cases, deaths, case.fatality.rates, start.date, end.date, death.undercount=1) {
  exposure.duration <- date.fxn(start.date, end.date)
  
  if(death.undercount>1) deaths <- round(deaths * death.undercount)
  deaths.data.wide <- create.data.wide(starting.ages, population, cases, deaths) 
  deaths.data.long <- create.data.long(deaths.data.wide)
  
  #estimate elasticity between [1] death counts and case fatality rates & [2] death counts and prevalence 
  fit <- logistf(deaths~log.case.fatality.rates+log.prevalence, data=deaths.data.long, weights=w, pl=T, firth=T)
  elasticity <- matrix(NA,nrow=2,ncol=2)
  rownames(elasticity) <- c("case.fatality.rates", "prevalence"); colnames(elasticity) <- c("pt.est", "se")
  elasticity["case.fatality.rates","pt.est"] <- summary(fit)$coef[2]
  elasticity["case.fatality.rates","se"] <- c(sqrt(diag(summary(fit)$var))[2])
  elasticity["prevalence","pt.est"] <- summary(fit)$coef[3]
  elasticity["prevalence","se"] <- c(sqrt(diag(summary(fit)$var))[3])
  
  #estimate age-group specific mortality rates using Firth's bias-reduced logistic regression
  fit.logistf <- logistf(deaths~as.factor(age), data=deaths.data.long, weights=w, pl=T, firth=T)
  pt.est <- fit.logistf$linear.predictors[1:nrow(deaths.data.wide)]
  sigma <- fit.logistf$var
  se <- rep(NA, length(pt.est))
  for (i in 1:length(se)) {
    C <- rep(0,length(se))
    C[i] <- 1 #assigns 1 to the ith age group 
    se[i] <- sqrt(t(C) %*% sigma %*% C)
  }
  mx.firth <- matrix(NA, nrow=length(starting.ages), ncol=3)
  rownames(mx.firth) <- starting.ages; colnames(mx.firth) <- c("pt.est", "ci, lower", "ci, upper")
  mx.firth[,"pt.est"] <- exp(pt.est)
  mx.firth[,"ci, lower"] <- exp(pt.est-qnorm(0.975)*se)
  mx.firth[,"ci, upper"] <- exp(pt.est+qnorm(0.975)*se)
  mx.firth <- (366/exposure.duration)*mx.firth
  
  fit.nb <- glm.nb(deaths~as.factor(starting.ages), data=deaths.data.wide, offset=log.exposure, control=glm.control(maxit=10^6))
  theta.par <- summary(fit.nb)$theta 
  #now fitting a glm with Neg Bin and an offset
  fit.nb2 <- glm(deaths~as.factor(starting.ages), data=deaths.data.wide, offset=log.exposure, family=neg.bin(theta.par))
  pred.nb <- predict(fit.nb2, newdata=data.frame(age=as.factor(seq(0,80,10)), log.exposure=deaths.data.wide$log.exposure), se.fit=TRUE)
  mx.nb <- exp(pred.nb$fit)/deaths.data.wide$exposure
  mx.nb <- (366/exposure.duration)*mx.nb
  
  #estimate age-group specific mortality rates using poisson regression
  fit.poisson <- glm(deaths~as.factor(age), data=deaths.data.wide, offset=log.exposure, family="poisson")
  pred.poisson <- predict(fit.poisson, newdata=data.frame(age=as.factor(seq(0,80,10)), log.exposure=deaths.data.wide$log.exposure), se.fit = TRUE)
  mx.poisson <- exp(pred.poisson$fit)/deaths.data.wide$exposure
  mx.poisson <- (366/exposure.duration)*mx.poisson
  
  return(list(deaths.data.wide=deaths.data.wide, elasticity=elasticity, mx.firth=mx.firth, mx.nb=mx.nb, mx.poisson=mx.poisson, fit.logistf=fit.logistf, fit.nb2=fit.nb2, fit.poisson=fit.poisson))
}

asdr.fxn <- function(population.standard, mx, elasticity, prevalence.observed, prevalence.scenario, deaths.data, asdr.scale=10^5) {
  relative.increase.prevalence <- (prevalence.scenario-prevalence.observed)/prevalence.observed
  asdr <- asdr.scale*sum(population.standard * ((1+elasticity*relative.increase.prevalence)*mx))
  return(asdr)
}

graph.mx.fxn <- function(mx) {
  ymin <- min(log10(100000*mx[,"ci, lower"]))
  ymax <- max(log10(100000*mx[,"ci, upper"]))
  cairo_pdf("covid19_mortality_rates_south_korea.pdf", width=8.5, height=5.5)
  par (mfrow=c(1,1), bg="white", mar=c(2,3,2,2), oma=c(2,2,2,2), font=1, cex=0.9, cex.main=0.9)
  plot(log10(100000*mx[,"pt.est"]), pch=19, ylim=c(ymin, ymax), xaxt="n", yaxt="n", ylab=NA, xlab=NA, bty="l", las=1)
  axis(1, at=1:9, paste(c("0-9", "10-19", "20-29", "30-39", "40-49", "50-59", "60-69", "70-79", "\u226580")))
  magaxis(side=2, unlog='y', las=1, minorn = 5, grid=TRUE, grid.lty=1, grid.col="grey90", grid.lwd =0.25)
  points(log10(100000*mx[,"pt.est"]), pch=19)
  for (i in 1:length(starting.ages)) lines(rep(i,2), log10(100000*mx[i,2:3]), lwd=1.5)
  mtext("Age Group (Years)", side=1, line=0.5, outer=TRUE, at=1/2, cex=1)
  mtext("Deaths per 100,000 Person-Years", side=2, line=0.5, outer=TRUE, at=1/2, cex=1)
  dev.off()
}

graph.alternative.mx.fxn <- function(mx.firth, mx.nb, mx.poisson) {
  ymin <- min(log10(c(100000*mx.firth[,"pt.est"], 100000*mx.nb, 100000*mx.poisson)))
  ymax <- max(log10(c(100000*mx.firth[,"pt.est"], 100000*mx.nb, 100000*mx.poisson)))
  cairo_pdf("mx.pdf", width=8.5, height=5.5)
  par (mfrow=c(1,1), bg="white", mar=c(2,3,2,2), oma=c(2,2,2,2), font=1, cex=0.9, cex.main=0.9)
  plot(log10(100000*mx.poisson), pch=19, ylim=c(ymin, ymax), xaxt="n", yaxt="n", ylab=NA, xlab=NA, bty="l", las=1, col=brewer.pal(3,"Set1")[3])
  lines(log10(100000*mx.poisson), col=brewer.pal(3,"Set1")[3])
  lines(log10(100000*mx.nb), col=brewer.pal(3,"Set1")[2])
  points(log10(100000*mx.nb), pch=19, col=brewer.pal(3,"Set1")[2])
  lines(log10(100000*mx.firth[,"pt.est"]), col=brewer.pal(3,"Set1")[1], lwd=2)
  points(log10(100000*mx.firth[,"pt.est"]), pch=19, col=brewer.pal(3,"Set1")[1])
  axis(1, at=1:9, paste(c("0-9", "10-19", "20-29", "30-39", "40-49", "50-59", "60-69", "70-79", "\u226580")))
  magaxis(side=2, unlog='y', las=1, minorn = 5, grid=TRUE, grid.lty=1, grid.col="grey90", grid.lwd =0.25)
  legend("bottomright", paste(c("Firth Weighted Logistic", "Negative Binomial", "Poisson")), bty="n", text.col=brewer.pal(3,"Set1"), col=brewer.pal(3,"Set1"), pch=19, lwd=c(2,1,1))
  mtext("Age Group (Years)", side=1, line=0.5, outer=TRUE, at=1/2, cex=1)
  mtext("Deaths per 100,000 Person-Years", side=2, line=0.5, outer=TRUE, at=1/2, cex=1)
  dev.off()
}

create.table.mx.fxn <- function(mx, mx.scale=10^5, age.label=c("0-9", "10-19", "20-29", "30-39", "40-49", "50-59", "60-69", "70-79", "\u226580")) {
  mx.table <- formatC(mx*mx.scale, digits=1, format="f", big.mark=",")
  rownames(mx.table) <- age.label
  return(mx.table)
}

create.table.elasticity.fxn <- function(elasticity) {
  pt.est <- elasticity[,1]
  se <- elasticity[,2]
  elasticity.table <- formatC(cbind(pt.est, se), digits=2, format="f")
  rownames(elasticity.table) <- c(simpleCap(str_c(lapply(rownames(elasticity), function(x) c(strsplit(x, "\\.")))[[1]][[1]], collapse=" ")),
                                  simpleCap(str_c(lapply(rownames(elasticity), function(x) c(strsplit(x, "\\.")))[[2]][[1]], collapse=" ")))
  return(elasticity.table)
}

simpleCap <- function(x) {
  s <- strsplit(x, " ")[[1]]
  paste(toupper(substring(s, 1,1)), substring(s, 2),
        sep="", collapse=" ")
}

asdr.mat.fxn <- function(prevalence.max, mortality.undercount.scenario, mortality.outcomes, population) {
  prevalence.observed <- sum(mortality.outcomes$deaths.data.wide$cases)/sum(population)
  prevalence.scenario <- seq(prevalence.observed, prevalence.max, length.out=100)
  asdr.mat <- matrix(NA, nrow=length(mortality.undercount.scenario), ncol=length(prevalence.scenario))
  rownames(asdr.mat) <- mortality.undercount.scenario
  colnames(asdr.mat) <- prevalence.scenario
  for (i in 1:length(mortality.undercount.scenario)) {
    for (j in 1:length(prevalence.scenario)) {
      asdr.mat[i,j] <- asdr.fxn(population.standard, mx[,1]*(1+mortality.undercount.scenario[i]), elasticity["prevalence","pt.est"], prevalence.observed, prevalence.scenario[j], deaths.data)
    }
  }
  return(asdr.mat)
}

graph.asdr <- function(asdr.mat) {
  prevalence.scenario <- as.numeric(dimnames(asdr.mat)[[2]])
  asdr.data <- data.frame(expand.grid(dimnames(asdr.mat)), c(asdr.mat))
  names(asdr.data) <- c("death.undercount", "prev", "asdr")
  asdr.levelplot <- 
    levelplot(asdr~death.undercount+prev, 
              data=asdr.data,
              labels=TRUE,
              at=seq(0,200,length.out=100),
              xlab="Scenario Death Undercount (%)",
              ylab="Scenario Prevalence (%)",
              col.regions = rev(colorRampPalette(brewer.pal(11,"RdYlBu"))(100)), colorkey = list(col = rev(colorRampPalette(brewer.pal(11, "RdYlBu"))(100))),
              scales=list(x=list(at=seq(1,100,length.out = 11), rot=0,labels=round(100*seq(0,1,length.out = 11),2)), 
                          y=list(at=seq(1,100,length.out = 11), rot=0,labels=round(100*seq(prevalence.scenario[1],last(prevalence.scenario), length.out = 11), 1))))
  cod <- rev(sort(cod))
  cod.labels <- paste(names(cod), c("(1st)", "(2nd)", "(3rd)", "(4th)", "(5th)", "(6th)", "(7th)", "(8th)", "(9th)", "(10th)"))
  cod.labels[cod.labels=="Heart Disease (2nd)"] <- "Heart Dis. (2nd)"
  cod.labels[cod.labels=="Cerebrovascular Disease (3rd)"] <- "Cerebrovascular Dis. (3rd)"
  cod.labels[cod.labels=="Intentional Self Harm (5th)"] <- "Intentional Self Harm (5th)"
  cod.labels[cod.labels=="Chronic Lower Resp. Tract Disease (7th)"] <- "CLRD (7th)"
  cod.labels[cod.labels=="Diabetes Mellitus (6th)"] <- "Diabetes (6th)"
  cod.labels[cod.labels=="Liver Disease (8th)"] <- "Liver Dis. (8th)"
  cod.labels[cod.labels=="Hypertensive Disease (10th)"] <- "Hypertensive Dis. (10th)"
  
  cairo_pdf("asdr.pdf", width=8.5, height=8.5)
  par (mfrow=c(1,1), bg="white", mar=c(2,3,2,2), oma=c(2,2,2,2), font=1, cex=0.9, cex.main=0.9)
  print(asdr.levelplot+
          as.layer(contourplot(asdr~death.undercount+prev, 
                               data=asdr.data,
                               col="black",
                               at=cod[-c(8,9)],
                               labels=list(labels=cod.labels[-c(8,9)], col="black", cex=0.75))))
  grid.text("COVID-19\nASDR", 0.95,0.97)
  dev.off()
}